import os
import re
import pickle

import numpy as np
import scipy.io as scio

from tqdm import tqdm

def process_csi_data(CSI_FOLDER='../WiFi/'):
    """
    处理CSI数据，提取起始时间戳。

    Args:
        CSI_FOLDER (str): 包含CSI数据的文件夹路径，默认为'../WiFi/'。
    Returns:
        CSI_START_TIME_LIST (list): 包含起始时间戳的列表，每个元素为 [CSI文件名, 起始时间戳]。

    # Example usage:
    # csi_start_time_list = process_csi_data(CSI_FOLDER='../WiFi/')
    """
    CSI_FOLDER_LIST = os.listdir(CSI_FOLDER)
    CSI_FOLDER_LIST = sorted(CSI_FOLDER_LIST, key=lambda x: int(re.findall(r'\d+', x)[0]))

    CSI_START_TIME_LIST = []

    for ONECSIDATAFILEPATH in tqdm(CSI_FOLDER_LIST):
        CSIDATANAME = 'rx_' + ONECSIDATAFILEPATH[0] + '_' + ONECSIDATAFILEPATH[1:7] + '_' + ONECSIDATAFILEPATH[7:13]

        csi_file = scio.loadmat(os.path.join(CSI_FOLDER, ONECSIDATAFILEPATH))

        csi_struct = csi_file[CSIDATANAME][0][0][0][0]

        START_TIME = np.float64(csi_struct[1][0][0][2][0] / 1000000000)

        CSI_START_TIME_LIST.append([ONECSIDATAFILEPATH[:13], START_TIME])
    
    return CSI_START_TIME_LIST

def CSI_PCA_DEMO_2(data, k):

    # k:前k个主成分
    data = (data - np.mean(data, axis=0)) / np.std(data, axis=0)  # 归一化数据
    pca = PCA()
    COEFF = pca.fit_transform(data)
    data_PCA = COEFF[:, :k]
 
    return data_PCA

def process_csi_data(WiFi_Mat_File, CSI_TIME):
    CSI_Mat = scio.loadmat(WiFi_Mat_File)
    CSIDATANAME = 'rx_' + CSI_TIME[0] + '_' + CSI_TIME[1:7] + '_' + CSI_TIME[7:13]

    csi_struct = CSI_Mat[CSIDATANAME][0][0][0][0]

    csi_data = csi_struct[3][0][0][13]  # 采集的CSI数据
    csi_tx1 = csi_data[:, 0:int(csi_data.shape[1] / 3)]
    csi_tx2 = csi_data[:, int(csi_data.shape[1] / 3):2 * int(csi_data.shape[1] / 3)]
    csi_tx3 = csi_data[:, 2 * int(csi_data.shape[1] / 3):]
    csidata = np.stack((csi_tx1, csi_tx2, csi_tx3), axis=2)

    csisystemtime = np.zeros([csi_data.shape[0]])

    for i in range(len(csi_struct[1][0][0][2])):
        csisystemtime[i] = np.float64(csi_struct[1][0][0][2][i] / 1000000000)

    Tx1 = csidata[:, :, 0]
    Tx2 = csidata[:, :, 1]
    ref_antenna = csidata[:, :, 2]
    ref_ante_conj = np.conj(ref_antenna)
    Tx2 = ref_ante_conj * Tx2

    CSIPCA2 = CSI_PCA_DEMO_2(abs(Tx2), 1)
    CSIPCA2 = np.squeeze(CSIPCA2)

    LowPassFitter = eng.LowPassFitter()
    # 进行滤波
    Tx2_LowPass = eng.hampel(eng.filter(LowPassFitter, matlab.double(abs(CSIPCA2).tolist())), 2000.0, 1.0)
    Tx2_LowPass = np.squeeze(np.array(Tx2_LowPass))

    time_offset = csisystemtime - csisystemtime[0]
    new_time = np.linspace(0, time_offset[-1], 150000)

    # 使用interp1d函数进行插值操作
    f = interp1d(time_offset, CSIPCA2.flatten(), kind='cubic')
    CSIPCA2new = f(new_time)

    csisystemtimenew = np.zeros([CSIPCA2new.shape[0]])

    for i in range(len(new_time)):
        csisystemtimenew[i] = new_time[i] + csisystemtime[0]

    return CSIPCA2new, csisystemtimenew

def csi_data_cut(video_start_timestamp, interv,CSI_timestamp,CSI_MAT,RADAR_ACTION_MAT_FILEPATH):
    
    CSI_ACTION_STARTTIME_LIST = []
    CSI_ACTION_STARTTIME_Index = []
    # 获取毫米波雷达和视频数据对应的时间戳
    for ii in range(len(CSI_timestamp)):
        if round(video_start_timestamp, 1) == round(CSI_timestamp[ii], 1):
            CSI_ACTION_STARTTIME_LIST.append(CSI_timestamp[ii])
            CSI_ACTION_STARTTIME_Index.append(ii)
            # 获取CSI中一个行为的开始时间和帧位
    if CSI_ACTION_STARTTIME_LIST:
        CSI_ACTION_STARTTIME = CSI_ACTION_STARTTIME_LIST[0]
        CSI_ACTION_FirstInex = CSI_ACTION_STARTTIME_Index[0]
        CSI_ACTION_EndInex = CSI_ACTION_FirstInex + interv
        ONEACTION_CSI_ORG_DATA = CSI_MAT[CSI_ACTION_FirstInex:CSI_ACTION_EndInex]


        scio.savemat(RADAR_ACTION_MAT_FILEPATH, {'RADAR_DATA': ONEACTION_CSI_ORG_DATA})


for ALL_Sensor_Time in SYNC_WIFI_FOLDER_LSIT:
    if os.path.exists(os.path.join(os.path.join(SYNC_WIFI_FOLDER
                                               ,ALL_Sensor_Time
                                               ,'WiFi'))):
        #读取CSI数据
        CSI_TIME=os.listdir(os.path.join(SYNC_WIFI_FOLDER
                                     ,ALL_Sensor_Time
                                     ,'WiFi'))[0]
        WiFi_Mat_File  = os.path.join(os.path.join(SYNC_WIFI_FOLDER
                                               ,ALL_Sensor_Time
                                               ,'WiFi')
                                               ,CSI_TIME)

        #读取视频数据时间戳
        VIDEO_TIME_DATA_FILEPATH = os.path.join(SYNC_WIFI_FOLDER, ALL_Sensor_Time, 'VIDEO', ALL_Sensor_Time, 'frame.pkl') #读取视频文件时间戳
        VIDEO_TIME_DATA = np.load(VIDEO_TIME_DATA_FILEPATH, allow_pickle=True)

        LABELED_CSI_MAT_time = os.path.join(LABELED_RADAR_MAT_FOLDER,ALL_Sensor_Time)
        #获取经过PCA、低通滤波、插值后的CSI矩阵、时间戳
        CSI_MAT,CSI_timestamp=process_csi_data(WiFi_Mat_File,CSI_TIME)

        time_action_list=all_action_index_list[ALL_Sensor_Time] #读取每隔2s的完整时间戳

        if not os.path.exists(LABELED_CSI_MAT_time):
            os.mkdir(LABELED_CSI_MAT_time)

        for action_name_index,action_index_list in tqdm(time_action_list.items()):
            for i in range(len(action_index_list)):
                RADAR_ACTION_MAT_FILEPATH = os.path.join(LABELED_CSI_MAT_time, str(action_name_index) + '_' + str(i) + '.mat')
                video_start_timestamp=VIDEO_TIME_DATA[action_index_list[i][0]].timestamp()

                csi_data_cut(video_start_timestamp,1000,CSI_timestamp,CSI_MAT,RADAR_ACTION_MAT_FILEPATH)